package com.crois.barrier.web.plugin.post;

import com.crois.barrier.common.enums.PluginEnum;
import com.crois.barrier.common.enums.PluginTypeEnum;
import com.crois.barrier.plugin.BarrierPlugin;
import com.crois.barrier.plugin.BarrierPluginChain;
import com.crois.barrier.web.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.netty.Connection;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

/**
 * @author Hou Ze Yu
 * @description netty 响应 处理 插件
 * @date 2020/7/13
 */
public class NettyClientResponsePlugin implements BarrierPlugin {

    private List<MediaType> streamingMediaTypes = Arrays
            .asList(MediaType.TEXT_EVENT_STREAM, MediaType.APPLICATION_STREAM_JSON);

    @Override
    public Mono<Void> execute(ServerWebExchange exchange, BarrierPluginChain chain) {
        return chain.execute(exchange).doOnError((throwable) -> {
         }).then(
                 Mono.defer(() -> {
                         Connection  connection = exchange.getAttribute(ServerWebExchangeUtils.BARRIER_CLIENT_RESPONSE_CONN_ATTR);
                         if (Objects.isNull(connection)){
                             return Mono.empty();
                         }
                         ServerHttpResponse response = exchange.getResponse();
                         NettyDataBufferFactory nettyDataBufferFactory = (NettyDataBufferFactory) response.bufferFactory();
                         Flux<NettyDataBuffer> nettyDataBufferFlux = connection
                                 .inbound()
                                 .receive()
                                 .retain()
                                 .map(nettyDataBufferFactory::wrap);
                         MediaType contentType = response.getHeaders().getContentType();
                         return isStreamMediaType(contentType)
                                 ? response.writeAndFlushWith(nettyDataBufferFlux.map(Flux::just))
                                 : response.writeWith(nettyDataBufferFlux);
        })).doOnCancel(() -> cleanup(exchange));
    }
    /**
     * 关闭连接
     * @param exchange
     */
    private void cleanup(ServerWebExchange exchange) {
           Connection connection = exchange.getAttribute(ServerWebExchangeUtils.BARRIER_CLIENT_RESPONSE_CONN_ATTR);
           if (connection!=null){
               connection.dispose();
           }
    }

    /**
     * 判断类型
     * @param contentType
     * @return
     */
    private boolean isStreamMediaType(MediaType contentType) {
        return contentType!=null&&streamingMediaTypes.stream().anyMatch(contentType::isCompatibleWith) ;
    }


    @Override
    public PluginTypeEnum pluginType() {
        return PluginTypeEnum.POST;
    }

    @Override
    public String name() {
        return PluginEnum.NETTY_CLIENT_RESPONSE.getName();
    }

    @Override
    public int order() {
        return PluginEnum.NETTY_CLIENT_RESPONSE.getCode();
    }

    @Override
    public Boolean skip(ServerWebExchange exchange) {
        return false;
    }
}
